"""
- reference: https://gitee.com/paddlepaddle/PaddleOCR
"""
import os
from tools import my_cv2_utils
from paddleocr import PaddleOCR, draw_ocr
from PIL import Image
import numpy as np
from bdtime import tt


# Paddleocr目前支持的多语言语种可以通过修改lang参数进行切换
# 例如`ch`, `en`, `fr`, `german`, `korean`, `japan`
ocr = PaddleOCR(use_angle_cls=True, lang="ch")  # need to run only once to download and load model into memory


def get_orc_result(img, coordinates=None, show=False, ret_one_line=True):
    if isinstance(img, str):
        assert os.path.exists(img)

    if coordinates:
        img = my_cv2_utils.crop_image(img, coordinates=coordinates)

    result = ocr.ocr(img, cls=False, det=False, rec=True)

    result = [i for i in result if i and i[0] and i[0][0]]

    if show:
        my_cv2_utils.show_img(img, resize=(600, 400))

    if ret_one_line:
        if result and result[0] and result[0][0]:
            text, score = result[0][0]
            result = [text, score]
        else:
            result = ['', 0]
    return result


if __name__ == '__main__':
    # img_path = '../data/ocr/none.png'  # None
    # img_path = '../data/ocr/a_5d9_1.png'  # None
    # img_path = '../data/ocr/a_5d9.png'
    # img_path = '../data/ocr/a_2d9.png'
    # img_path = '../data/ocr/a_8d9.png'
    # img_path = '../data/ocr/a_8d9_1.png'
    # img_path = '../data/ocr/a_complete.png'
    img_path = '../data/ocr/screen__0d9.jpg'

    # result = get_orc_result(img_path, show=True)
    _img = my_cv2_utils.read_image(img_path)
    print('_img.shape:', _img.shape)
    # my_cv2_utils.show_img(_img)

    # _img = my_cv2_utils.crop_image(img, coordinates=[(0, 0), (100, 100)])
    # img = my_cv2_utils.crop_image(_img, coordinates=[(25, 105), (90, 115)])

    coordinates = [(65, 104), (82, 115)]
    img = my_cv2_utils.crop_image(_img, coordinates=coordinates)
    # my_cv2_utils.show_img(img)
    print('img.shape:', img.shape)

    result = get_orc_result(img)
    print('--- ocr_result:', result)

    times = 20
    tt.__init__()
    for i in range(times):
        result = get_orc_result(img_path)
    print(f'--- mean cost time: {tt.now() / times: .3f}')
